import argparse
import os.path as osp

from torch.utils.data import DataLoader
import torchvision
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
import torch
from tqdm import tqdm

from lenet import LeNet
from resnet18 import ResNet
from new_mnist import NewMNIST
from new_cifar10 import NewCIFAR10


def prepare(task, dataset_root):
    if task == "mnist":
        model = LeNet(10)
        model.cuda()

        train_ds = NewMNIST(
            dataset_root, train=True,
            transform=transforms.ToTensor()
        )
        test_ds = NewMNIST(
            dataset_root, train=False,
            transform=transforms.ToTensor()
        )
    elif task == "cifar10":
        model = ResNet(11)
        model.cuda()
        dataset_root, extra_imgs_root = dataset_root.split(';')
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        train_ds = NewCIFAR10(
            dataset_root, extra_imgs_root,
            train=True, transform=transform
        )
        test_ds = NewCIFAR10(
            dataset_root, extra_imgs_root,
            train=False, transform=transform
        )

    trainloader = DataLoader(
        train_ds, batch_size=64,
        shuffle=True
    )
    testloader = DataLoader(
        test_ds, batch_size=64,
        shuffle=False
    )
    return model, trainloader, testloader


def main():
    parser = argparse.ArgumentParser("train model")
    parser.add_argument('--task', type=str, help='task type',
                        default='mnist', choices=['mnist', 'cifar10'])
    parser.add_argument('--dataset-root', type=str)
    parser.add_argument('--ckpt-path', type=str,
                        help='checkpoint path')
    parser.add_argument('--epochs', type=int, default=20, metavar='N',
                        help='number of epochs to train')
    args = parser.parse_args()

    model, trainloader, testloader = prepare(args.task, args.dataset_root)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    for epoch in range(args.epochs):
        model.train()
        running_loss = 0.0
        for i, data in enumerate(tqdm(trainloader)):
            inputs, labels = data
            inputs, labels = inputs.to("cuda"), labels.to("cuda")

            optimizer.zero_grad()

            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        print(f'[epoch #{epoch + 1}] loss: {running_loss / len(trainloader)}')
        torch.save(
            model.state_dict(),
            osp.join(args.ckpt_path, f"{epoch}.pth")
        )

        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for i, data in enumerate(testloader):
                inputs, labels = data
                inputs, labels = inputs.to("cuda"), labels.to("cuda")

                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        print(f'[epoch #{epoch+ 1}] accuracy: {correct / total}')


if __name__ == "__main__":
    main()
